/*
 * LD_PRELOAD hook to intercept file deletion calls
 * 编译: gcc -shared -fPIC -o capture_hivm.so capture_hivm_ldpreload.c -ldl
 * 使用: LD_PRELOAD=./capture_hivm.so bishengir-compile ...
 */
#define _GNU_SOURCE
#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <time.h>
#include <stdarg.h>

#define SAVE_DIR "/tmp/hivm_captures"
#define MAX_PATH 4096

// 获取当前时间戳目录名
static void get_timestamp_dir(char *dir, size_t dir_size) {
    time_t now = time(NULL);
    struct tm *tm_info = localtime(&now);
    snprintf(dir, dir_size, "%s/%04d%02d%02d_%02d%02d%02d", 
             SAVE_DIR,
             tm_info->tm_year + 1900,
             tm_info->tm_mon + 1,
             tm_info->tm_mday,
             tm_info->tm_hour,
             tm_info->tm_min,
             tm_info->tm_sec);
}

// 保存文件到安全位置
static void save_file_before_delete(const char *path) {
    // 捕获所有bishengir-compile目录下的文件
    if (strstr(path, "bishengir-compile") == NULL) {
        return;  // 只处理bishengir-compile目录下的文件
    }
    
    // 排除一些不需要的文件
    if (strstr(path, ".swp") != NULL || strstr(path, ".tmp") != NULL) {
        return;
    }
    
    // 检查文件是否存在
    struct stat st;
    if (stat(path, &st) != 0) {
        return;
    }
    
    // 创建时间戳目录
    char timestamp_dir[MAX_PATH];
    get_timestamp_dir(timestamp_dir, sizeof(timestamp_dir));
    char cmd[MAX_PATH + 20];
    snprintf(cmd, sizeof(cmd), "mkdir -p %s", timestamp_dir);
    system(cmd);
    
    // 生成保存路径
    char save_path[MAX_PATH];
    const char *basename = strrchr(path, '/');
    if (basename == NULL) {
        basename = path;
    } else {
        basename++;
    }
    
    // 获取父目录名
    char parent_dir[256] = {0};
    const char *last_slash = strrchr(path, '/');
    if (last_slash != NULL) {
        const char *prev_slash = strrchr(path, '/');
        if (prev_slash != NULL && prev_slash != last_slash) {
            // 找到倒数第二个斜杠
            const char *p = path;
            const char *second_last = NULL;
            while ((p = strchr(p + 1, '/')) != NULL && p < last_slash) {
                second_last = p;
            }
            if (second_last != NULL) {
                int len = last_slash - second_last - 1;
                if (len > 0 && len < sizeof(parent_dir) - 1) {
                    strncpy(parent_dir, second_last + 1, len);
                    parent_dir[len] = '\0';
                }
            }
        }
    }
    
    // 构建保存路径，保持相对路径结构
    // 例如: /tmp/bishengir-compile-xxx/file.mlir -> timestamp_dir/bishengir-compile-xxx/file.mlir
    const char *bishengir_pos = strstr(path, "bishengir-compile");
    if (bishengir_pos != NULL) {
        // 提取bishengir-compile-xxx之后的部分
        const char *relative_path = bishengir_pos;
        snprintf(save_path, sizeof(save_path), "%s/%s", timestamp_dir, relative_path);
        // 确保父目录存在
        char *last_slash = strrchr(save_path, '/');
        if (last_slash != NULL) {
            *last_slash = '\0';
            snprintf(cmd, sizeof(cmd), "mkdir -p %s", save_path);
            system(cmd);
            *last_slash = '/';
        }
    } else {
        // 回退方案：使用basename
        snprintf(save_path, sizeof(save_path), "%s/%s", timestamp_dir, basename);
    }
    
    // 复制文件
    FILE *src = fopen(path, "rb");
    if (src != NULL) {
        FILE *dst = fopen(save_path, "wb");
        if (dst != NULL) {
            char buffer[8192];
            size_t n;
            while ((n = fread(buffer, 1, sizeof(buffer), src)) > 0) {
                fwrite(buffer, 1, n, dst);
            }
            fclose(dst);
            fprintf(stderr, "[HIVM Hook] 已保存文件: %s -> %s (大小: %ld bytes)\n", path, save_path, (long)st.st_size);
        }
        fclose(src);
    }
}

// Hook unlink
int unlink(const char *pathname) {
    int (*original_unlink)(const char *) = dlsym(RTLD_NEXT, "unlink");
    
    save_file_before_delete(pathname);
    
    return original_unlink(pathname);
}

// Hook unlinkat
int unlinkat(int dirfd, const char *pathname, int flags) {
    int (*original_unlinkat)(int, const char *, int) = dlsym(RTLD_NEXT, "unlinkat");
    
    // 构建完整路径并保存所有bishengir-compile目录下的文件
    if (pathname != NULL && strstr(pathname, "bishengir-compile") != NULL) {
        char full_path[MAX_PATH];
        if (pathname[0] == '/') {
            strncpy(full_path, pathname, sizeof(full_path) - 1);
            full_path[sizeof(full_path) - 1] = '\0';
        } else {
            // 尝试从dirfd获取路径
            char fd_path[MAX_PATH];
            snprintf(fd_path, sizeof(fd_path), "/proc/self/fd/%d", dirfd);
            char link_target[MAX_PATH];
            ssize_t len = readlink(fd_path, link_target, sizeof(link_target) - 1);
            if (len > 0) {
                link_target[len] = '\0';
                snprintf(full_path, sizeof(full_path), "%s/%s", link_target, pathname);
            } else {
                // 回退方案：假设在/tmp下
                snprintf(full_path, sizeof(full_path), "/tmp/%s", pathname);
            }
        }
        save_file_before_delete(full_path);
    }
    
    return original_unlinkat(dirfd, pathname, flags);
}

// Hook remove
int remove(const char *pathname) {
    int (*original_remove)(const char *) = dlsym(RTLD_NEXT, "remove");

    save_file_before_delete(pathname);

    return original_remove(pathname);
}

// 保存已创建的文件
static void save_created_file(const char *path) {
    // 捕获所有/tmp下的编译相关文件
    if (strncmp(path, "/tmp/", 5) != 0) {
        return;
    }

    // 需要捕获的文件类型
    int should_capture = 0;

    // bishengir相关目录
    if (strstr(path, "bishengir-compile") != NULL ||
        strstr(path, "bishengir-bc-manager") != NULL) {
        should_capture = 1;
    }

    // 特定扩展名
    if (strstr(path, ".mlir") != NULL ||
        strstr(path, ".ll") != NULL ||
        strstr(path, ".bc") != NULL ||
        strstr(path, ".o") != NULL ||
        strstr(path, ".c") != NULL ||
        strstr(path, ".out") != NULL ||
        strstr(path, ".so") != NULL) {
        should_capture = 1;
    }

    if (!should_capture) {
        return;
    }

    // 排除一些不需要的文件
    if (strstr(path, ".swp") != NULL) {
        return;
    }

    // 创建时间戳目录
    char timestamp_dir[MAX_PATH];
    get_timestamp_dir(timestamp_dir, sizeof(timestamp_dir));
    char cmd[MAX_PATH + 20];
    snprintf(cmd, sizeof(cmd), "mkdir -p %s", timestamp_dir);
    system(cmd);

    // 构建保存路径，保持相对路径结构
    char save_path[MAX_PATH];
    const char *relative = path + 5; // skip "/tmp/"
    snprintf(save_path, sizeof(save_path), "%s/tmp/%s", timestamp_dir, relative);

    // 确保父目录存在
    char *last_slash = strrchr(save_path, '/');
    if (last_slash != NULL) {
        *last_slash = '\0';
        snprintf(cmd, sizeof(cmd), "mkdir -p %s", save_path);
        system(cmd);
        *last_slash = '/';
    }

    // 复制文件
    struct stat st;
    if (stat(path, &st) == 0 && S_ISREG(st.st_mode)) {
        FILE *src = fopen(path, "rb");
        if (src != NULL) {
            FILE *dst = fopen(save_path, "wb");
            if (dst != NULL) {
                char buffer[8192];
                size_t n;
                while ((n = fread(buffer, 1, sizeof(buffer), src)) > 0) {
                    fwrite(buffer, 1, n, dst);
                }
                fclose(dst);
                fprintf(stderr, "[HIVM Hook] 已捕获文件: %s -> %s (大小: %ld bytes)\n",
                        path, save_path, (long)st.st_size);
            }
            fclose(src);
        }
    }
}

// 跟踪已处理的文件描述符和路径
typedef struct {
    int fd;
    char path[MAX_PATH];
    int needs_capture;
} fd_info_t;

#define MAX_FDS 1024
static fd_info_t fd_table[MAX_FDS];
static int fd_table_initialized = 0;

static void init_fd_table() {
    if (!fd_table_initialized) {
        memset(fd_table, 0, sizeof(fd_table));
        for (int i = 0; i < MAX_FDS; i++) {
            fd_table[i].fd = -1;
        }
        fd_table_initialized = 1;
    }
}

static void register_fd(int fd, const char *path, int needs_capture) {
    init_fd_table();
    if (fd >= 0 && fd < MAX_FDS) {
        fd_table[fd].fd = fd;
        strncpy(fd_table[fd].path, path, MAX_PATH - 1);
        fd_table[fd].path[MAX_PATH - 1] = '\0';
        fd_table[fd].needs_capture = needs_capture;
    }
}

static void unregister_fd(int fd) {
    init_fd_table();
    if (fd >= 0 && fd < MAX_FDS) {
        // 在关闭时捕获文件
        if (fd_table[fd].needs_capture && fd_table[fd].path[0] != '\0') {
            save_created_file(fd_table[fd].path);
        }
        fd_table[fd].fd = -1;
        fd_table[fd].path[0] = '\0';
        fd_table[fd].needs_capture = 0;
    }
}

// Hook open
int open(const char *pathname, int flags, ...) {
    static int (*original_open)(const char *, int, ...) = NULL;
    if (!original_open) {
        original_open = dlsym(RTLD_NEXT, "open");
    }

    mode_t mode = 0;
    if (flags & O_CREAT) {
        va_list args;
        va_start(args, flags);
        mode = va_arg(args, mode_t);
        va_end(args);
    }

    int fd = original_open(pathname, flags, mode);

    // 如果是创建文件，记录下来
    if (fd >= 0 && (flags & O_CREAT) && pathname != NULL) {
        int needs_capture = 0;
        if (strncmp(pathname, "/tmp/", 5) == 0) {
            if (strstr(pathname, "bishengir") != NULL ||
                strstr(pathname, ".mlir") != NULL ||
                strstr(pathname, ".ll") != NULL ||
                strstr(pathname, ".bc") != NULL ||
                strstr(pathname, ".o") != NULL ||
                strstr(pathname, ".c") != NULL ||
                strstr(pathname, ".out") != NULL ||
                strstr(pathname, ".so") != NULL) {
                needs_capture = 1;
            }
        }
        register_fd(fd, pathname, needs_capture);
    }

    return fd;
}

// Hook openat
int openat(int dirfd, const char *pathname, int flags, ...) {
    static int (*original_openat)(int, const char *, int, ...) = NULL;
    if (!original_openat) {
        original_openat = dlsym(RTLD_NEXT, "openat");
    }

    mode_t mode = 0;
    if (flags & O_CREAT) {
        va_list args;
        va_start(args, flags);
        mode = va_arg(args, mode_t);
        va_end(args);
    }

    int fd = original_openat(dirfd, pathname, flags, mode);

    // 如果是创建文件，记录下来
    if (fd >= 0 && (flags & O_CREAT) && pathname != NULL) {
        char full_path[MAX_PATH];
        if (pathname[0] == '/') {
            strncpy(full_path, pathname, sizeof(full_path) - 1);
            full_path[sizeof(full_path) - 1] = '\0';
        } else {
            char fd_path[MAX_PATH];
            snprintf(fd_path, sizeof(fd_path), "/proc/self/fd/%d", dirfd);
            char link_target[MAX_PATH];
            ssize_t len = readlink(fd_path, link_target, sizeof(link_target) - 1);
            if (len > 0) {
                link_target[len] = '\0';
                snprintf(full_path, sizeof(full_path), "%s/%s", link_target, pathname);
            } else {
                strncpy(full_path, pathname, sizeof(full_path) - 1);
                full_path[sizeof(full_path) - 1] = '\0';
            }
        }

        int needs_capture = 0;
        if (strncmp(full_path, "/tmp/", 5) == 0) {
            if (strstr(full_path, "bishengir") != NULL ||
                strstr(full_path, ".mlir") != NULL ||
                strstr(full_path, ".ll") != NULL ||
                strstr(full_path, ".bc") != NULL ||
                strstr(full_path, ".o") != NULL ||
                strstr(full_path, ".c") != NULL ||
                strstr(full_path, ".out") != NULL ||
                strstr(full_path, ".so") != NULL) {
                needs_capture = 1;
            }
        }
        register_fd(fd, full_path, needs_capture);
    }

    return fd;
}

// Hook close
int close(int fd) {
    static int (*original_close)(int) = NULL;
    if (!original_close) {
        original_close = dlsym(RTLD_NEXT, "close");
    }

    unregister_fd(fd);

    return original_close(fd);
}

// Hook creat
int creat(const char *pathname, mode_t mode) {
    static int (*original_creat)(const char *, mode_t) = NULL;
    if (!original_creat) {
        original_creat = dlsym(RTLD_NEXT, "creat");
    }

    int fd = original_creat(pathname, mode);

    if (fd >= 0 && pathname != NULL) {
        int needs_capture = 0;
        if (strncmp(pathname, "/tmp/", 5) == 0) {
            if (strstr(pathname, "bishengir") != NULL ||
                strstr(pathname, ".mlir") != NULL ||
                strstr(pathname, ".ll") != NULL ||
                strstr(pathname, ".bc") != NULL ||
                strstr(pathname, ".o") != NULL ||
                strstr(pathname, ".c") != NULL ||
                strstr(pathname, ".out") != NULL ||
                strstr(pathname, ".so") != NULL) {
                needs_capture = 1;
            }
        }
        register_fd(fd, pathname, needs_capture);
    }

    return fd;
}
